import json
from tqdm import tqdm


'''
alpaca 数据集格式：
    [ 
        { "instruction":"", "input":"", "output":"" },
        { "instruction":"", "input":"", "output":"" },
    ]
'''
def convert_to_alpaca_json(file_path):
    dataset = []
    # 打开文件，指定使用GBK编码
    with open(file_path, 'r',encoding='utf-8') as file:
        # 遍历文件的每一行
        for line in file:
            try:
                # 使用Tab分割每一行的内容
                parts = line.strip().split('\t')
                # 确保分割后有两部分
                if len(parts) == 2:
                    # 打印分割后的两段内容
                    # print(f"user:{parts[0]}, \nagent:{parts[1]}")
                    entry = {
                        "output": parts[1],
                        "input": "",  # 根据需要添加输入字段
                        "instruction": parts[0]
                    }
                    dataset.append(entry)
                else:
                    # 如果没有两部分，则跳过该行
                    print("Warning: Line does not contain two parts. Skipping line:")
                    # print(line.strip())
            except Exception as e:
                # 打印错误信息并跳过当前行
                print(f"Error processing line: {line.strip()}. Error: {e}")
    return dataset



'''
多轮对话数据集格式：
    [
        {"id": "x3959", "conversations": [{"from": "user", "value": "xxx"}, {"from": "assistant", "value": "xxx"}, {"from": "user", "value": "xxx"}, {"from": "assistant", "value": "xxx"}]},
        {"id": "x3960", "conversations": [{"from": "user", "value": "xxx"}, {"from": "assistant", "value": "xxx"}, {"from": "user", "value": "xxx"}, {"from": "assistant", "value": "xxx"}]},
    ]
'''
def convert_to_json(file_path):
    conversations = []
    with open(file_path, 'r', encoding='utf-8') as file:
        lines = file.readlines()
        total_iterations = len(lines)
        current_id = 1
        current_conversation = []
        for line in tqdm(lines, desc="Processing conversations"):
            line = line.strip()
            if line:
                current_conversation.append({"from": "user" if len(current_conversation) % 2 == 0 else "assistant", "value": line})
            else:
                if current_conversation:
                    conversations.append({"id": f"x{current_id}", "conversations": current_conversation})
                    current_conversation = []
                    current_id += 1   
        # Append the last conversation if any
        if current_conversation:
            conversations.append({"id": f"x{current_id}", "conversations": current_conversation})
    return conversations



'''
多轮对话保存为问答：
    [ 
        { "instruction":"", "input":"", "output":"" }, ··· 
    ]
举例:
    (原始对话为两段)
    捏的这么可爱
    眼睛特别搞笑这土也不好捏但就是觉得挺可爱
    特别可爱啊

    今天好点了吗？
    一天比一天严重
    吃药不管用，去打一针。别拖着
转换后：
    [ 
        { "instruction":"捏的这么可爱", "input":"", "output":"眼睛特别搞笑这土也不好捏但就是觉得挺可爱" }, { "instruction":"眼睛特别搞笑这土也不好捏但就是觉得挺可爱", "input":"", "output":"特别可爱啊" },
        { "instruction":"今天好点了吗？", "input":"", "output":"一天比一天严重" }, { "instruction":"一天比一天严重", "input":"", "output":"吃药不管用，去打一针。别拖着" },
    ]
'''
def process_dialogues(file_path):
    conversations = []
    with open(file_path, 'r', encoding='utf-8') as file:
        lines = file.readlines()
        current_instruction = None

        for line in tqdm(lines, desc="Processing conversations"):
            line = line.strip()
            if line:
                if current_instruction is None:
                    current_instruction = line
                else:
                    conversations.append({
                        "instruction": current_instruction,
                        "input": "",
                        "output": line
                    })
                    current_instruction = line
            else:
                current_instruction = None
        # Append the last dialogue if any
        if current_instruction:
            conversations.append({
                "instruction": current_instruction,
                "input": "",
                "output": ""
            })
    return conversations




if __name__ == '__main__':
    # Specify the path to your train.txt file
    file_path = 'datasets\GPT2-chitchat.txt'
    save_path = 'GPT2-chitchat_single.json'

    # Process dialogues
    # 选择使用哪个函数，则设置为True
    json_data = None
    if False: json_data = convert_to_alpaca_json(file_path)
    if False: json_data = convert_to_json(file_path)
    if False: json_data = process_dialogues(file_path)

    # Output the JSON data
    with open(save_path, 'w', encoding='utf-8') as file:
        json.dump(json_data, file, ensure_ascii=False, indent=2)
    print("done.")
